'''
Created on Oct 5, 2009

@author: aamn
'''
import numpy as np
import numpy.ma as ma


def conv_waves(wave, unit=['angstroms', 'angstroms']):
    """ convert between wave units: 
    [Angstroms, nm, micron, wavenumber,Kev, etc]
    """
    ck1 = ['hz', 'wn', 'kev', 'angstrom', 'nm', 'micron', 'mm', 'cm', 'm']
    c = np.double(2.99792458e18)
    ev = np.double(12.3984282)
    utoA = [lambda u: c / u,
            lambda u: (u ** (-1)) * np.double(1e7),
            lambda u: (u / ev) ** (-1),
            lambda u: u,
            lambda u: u * np.double(1e1),
            lambda u: u * np.double(1e4),
            lambda u: u * np.double(1e7),
            lambda u: u * np.double(1e8),
            lambda u: u * np.double(1e10),
            ]
    Atou = [lambda a: c / a,
            lambda a: (a / np.double(1e7)) ** (-1),
            lambda a: ev * (a ** (-1)),
            lambda a: a,
            lambda a: a / np.double(1e1),
            lambda a: a / np.double(1e4),
            lambda a: a / np.double(1e7),
            lambda a: a / np.double(1e8),
            lambda a: a / np.double(1e10),
            ]

    for i, v in enumerate(unit):
        if v.endswith('s'): unit[i] = v.rstrip('s').lower()
        else: unit[i] = v.lower()

    if unit[0] == unit[1]: return wave
    if unit[0] in ck1 and unit[1] in ck1:
        return Atou[ck1.index(unit[1])](utoA[ck1.index(unit[0])](wave))
    else: raise  ValueError('unknown units for conversion')

def parse_wtype(wtype):
    """ parse the FITS header wtype keys
    
    @param wtype:    string of space delimited key=value pairs.
    """
    return dict([tuple(foo.split('=')) for foo in wtype.split(' ') ])

def eqw(indice, wave, flux):
    """ EqW generator.  This is a stub.

    @param indice:  A spectral index (dictionary) 
    @param wave:    masked wavelength vector
    @param flux:    masked flux vector
    """
    # 1. from 2 continuum points on either side of a line
    #    calculate the "continuum" flux from a linear interpolation
    #    across the line. This is a small flux vector with same # elements
    #    as in the line band.

    # 2. Do a "Trapezoidal" sum across the line for:
    #    1. continuum - flux 
    #    Actually what is subtracted from which is important.
    #    Kevin subtracts the line from the continuum.
    #    absorbed_flux = TSUM(lambda[inline], continuum - flux[inline]) [IDL]
    #    absorbed_flux = np.trapz(continuum - flux[inband], wave[inband]) [py]
    #    2. flux divided by the total bandwidth
    #    I don't understand this; each flux pixel is divided by the inbandwidth
    # 3. Divide these fluxes by the interpolated continuum at the line center
    # 
    return None

def specbandwmean(bands, wave, flux):
    """ For a set of spectral intervals or "bands" return the weighted mean 
    and standard deviation of the band and the total number of spectral 
    elements in the bands.
    
    @param bands:   list containing band elements. each band element is a 
                    3 element list: [min wave, max wave, weight (<1)]
                    The sum of the weights in all the bands should equal 1. 
    @param wave:    wavelength vector (masked ndarray)
    @param flux:    flux vector (masked ndarray)
    """
    wmean = np.sum([band[2] * (ma.mean(flux[(wave > band[0]) & (wave < band[1])]))
                  for band in bands
                  if ma.masked_outside(wave, band[0], band[1]).count() > 0])
    welem = np.sum([ma.masked_outside(wave, band[0], band[1]).count()
                  for band in bands
                  if ma.masked_outside(wave, band[0], band[1]).count() > 0])
    if welem > 0:
      snoise = max([ma.std(flux[(wave > band[0]) & (wave < band[1])])
                  for band in bands
                  if ma.masked_outside(wave, band[0], band[1]).count() > 0])
    else:
      snoise = np.nan
    #
    return wmean, welem, snoise

def specindex(indice, wave, flux):
    """ For a spectral index, specified by a set of spectral intervals inside
    (N) and outside (D) of the feature, return the the ratio of N/D, the 
    difference of the index and a noise adjusted N/D, and the SNR in the 
    sideband. 
    
    @param indice:  A spectral index (dictionary) 
    @param wave:    masked wavelength vector
    @param flux:    masked flux vector
    """
    lineband, lineelem, linenoise = specbandwmean(indice['N'], wave, flux)
    sideband, sideelem, sidenoise = specbandwmean(indice['D'], wave, flux)
    sidesnr = sideband / sidenoise
    sidevardN = sidenoise * sidenoise / sideelem
    specindex = lineband / sideband
    indexdiff = abs(
                  specindex -
                  (lineband / (sideband + np.sqrt(sidevardN)))
                  )
    return specindex, indexdiff, sidesnr

def specsum(wave):
    """ return list of useful quantities for a vector of spectral units
    believes it is getting a numpy array of some kind
    """
    outs = ('min', 'max', 'ptp', 'length', 'dx', 'kind', 'dtype')
    dwave = np.zeros(wave.shape, dtype=wave.dtype)
    dwave[:-1] = (wave[1:] - wave[:-1])
    dwave[-1] = dwave[-2]
    if len(np.unique(dwave)) != 1:
        kind = 'nonlinear'
    else:
        kind = 'linear'
    outd = (wave.min(), wave.max(), wave.ptp(), wave.shape[0], dwave, kind,
            wave.dtype)
    return dict(zip(outs, outd))

def vac2air(wave):
    """ straight idl->python dump of vactoair.pro
    see http://idlastro.gsfc.nasa.gov/ftp/pro/astro/vactoair.pro
    """
    coef = np.array([1.0, 2.735182e-4, 131.4182, 2.76249e8], dtype=np.double)
    fact = coef[0] + \
           coef[1] + \
           coef[2] / (wave * wave) + \
           coef[3] / (wave * wave * wave * wave)
    fact = fact * (wave >= 2000.) + 1.0 * (wave < 2000.0)
    #
    return wave / fact
